#!/usr/bin/env python3
"""
ESP32 自动刷写工具
使用esptool来刷写编译好的固件到ESP32
"""

import sys
import os
import subprocess
import time
import argparse

def run_command(cmd, description=""):
    """执行命令并显示结果"""
    print(f"🔄 {description}")
    print(f"💻 执行: {cmd}")
    
    try:
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        if result.returncode == 0:
            print(f"✅ 成功: {description}")
            if result.stdout:
                print(f"📤 输出:\n{result.stdout}")
            return True
        else:
            print(f"❌ 失败: {description}")
            if result.stderr:
                print(f"⚠️ 错误:\n{result.stderr}")
            return False
    except Exception as e:
        print(f"❌ 执行失败: {e}")
        return False

def detect_esp32(port="/dev/ttyUSB0"):
    """检测ESP32设备"""
    cmd = f"esptool.py --port {port} chip_id"
    return run_command(cmd, f"检测ESP32设备 ({port})")

def erase_flash(port="/dev/ttyUSB0"):
    """擦除Flash"""
    cmd = f"esptool.py --port {port} erase_flash"
    return run_command(cmd, "擦除ESP32 Flash")

def flash_firmware(port="/dev/ttyUSB0", firmware_path=""):
    """刷写固件"""
    if not firmware_path or not os.path.exists(firmware_path):
        print(f"❌ 固件文件不存在: {firmware_path}")
        return False
    
    cmd = f"esptool.py --port {port} --baud 460800 write_flash 0x1000 {firmware_path}"
    return run_command(cmd, f"刷写固件 ({firmware_path})")

def create_simple_firmware():
    """创建一个简单的二进制固件文件"""
    print("🔄 创建简单测试固件...")
    
    # 这是一个基本的ESP32固件头
    # 实际应用中需要Arduino IDE编译完整固件
    firmware_data = bytearray([
        0xE9, 0x01, 0x02, 0x03,  # ESP32 Magic Number + header
        0x00, 0x00, 0x00, 0x00,  # Entry point
    ] + [0xFF] * 1016)  # Padding
    
    with open("simple_firmware.bin", "wb") as f:
        f.write(firmware_data)
    
    print("✅ 创建了简单固件文件: simple_firmware.bin")
    return "simple_firmware.bin"

def main():
    parser = argparse.ArgumentParser(description="ESP32 刷写工具")
    parser.add_argument("--port", default="/dev/ttyUSB0", help="ESP32串口设备")
    parser.add_argument("--firmware", help="固件文件路径")
    parser.add_argument("--erase", action="store_true", help="擦除Flash")
    parser.add_argument("--detect", action="store_true", help="仅检测设备")
    
    args = parser.parse_args()
    
    print("=== ESP32 刷写工具 ===")
    print(f"设备端口: {args.port}")
    
    # 激活虚拟环境中的esptool
    os.environ["PATH"] = "/home/mobox/hardware/esp_venv/bin:" + os.environ.get("PATH", "")
    
    # 检测设备
    if not detect_esp32(args.port):
        print("❌ 无法检测到ESP32设备，请检查连接")
        return 1
    
    if args.detect:
        print("✅ ESP32设备检测完成")
        return 0
    
    # 擦除Flash
    if args.erase:
        if not erase_flash(args.port):
            print("❌ 擦除Flash失败")
            return 1
    
    # 刷写固件
    if args.firmware:
        if not flash_firmware(args.port, args.firmware):
            print("❌ 刷写固件失败")
            return 1
    else:
        print("ℹ️ 没有指定固件文件")
        print("💡 提示: 使用Arduino IDE编译ESP32固件，然后:")
        print(f"    python3 {sys.argv[0]} --firmware your_firmware.bin")
    
    print("🎉 操作完成!")
    return 0

if __name__ == "__main__":
    sys.exit(main()) 